package dataStructure.tree.huffmanTree;

import java.sql.SQLOutput;
import java.util.*;

/**
 * @author lhn
 * @date 2024-09-05 17:05
 */

public class HuffmanCode {
    public static void main(String[] args) {
        String content = "i like like like java do you like a java";
        byte[] bytes = content.getBytes();  //字符串转字节数组

        List<NodePlus> nodes = getNodes(bytes);
        System.out.println(nodes);

        NodePlus root = huffmanCoding(nodes);
        System.out.println("前序遍历");
        preOrder(root);

        //根据赫夫曼树码，生成赫夫曼树对应的赫夫曼编码表
        getCodes(root);
        System.out.println("赫夫曼编码表:");
        System.out.println(huffmanCodes);
        //{32=01, 97=100, 100=11000, 117=11001, 101=1110, 118=11011, 105=101, 121=11010, 106=0010, 107=1111, 108=000, 111=0011}

    }

    //1.将赫夫曼编码表存放在map中Map<Byte,String>
    static Map<Byte, String > huffmanCodes = new HashMap<>();
    //2.在生成赫夫曼编码表时，需要去拼接路径，定义一个StringBuilder存储叶子节点的路径
    static StringBuilder stringBuilder = new StringBuilder();

    /**
    * 将字符串对应的byte[] 数组，通过生成的huffman编码表，生成赫夫曼编码对应的字符串
     * @param bytes 原始字符串对应的byte[]
     * @param huffmanCodes 生成的赫夫曼编码表
     * @return bytes 使用huffman编码表，进行压缩后的字节数组 例：10010001()
     * 8位对应一个byte huffmanCodeBytes[0] = 10101000 => byte
    */
    //TODO
//    public static byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes) {
//
//    }

    /**
    * 重载
     * @param root 传入根节点
     * @return huffman编码表
    */
    public static Map<Byte, String> getCodes(NodePlus root) {
        if (root == null) {
            return null;
        }
        getCodes(root.getLeft(), "0", stringBuilder);   //左子树
        getCodes(root.getRight(), "1", stringBuilder);  //右子树
        return huffmanCodes;
    }

    /**
    * 将传入的node节点的所有叶子节点的赫夫曼编码，保存在huffmanCodes集合中
     * @param node 传入节点
     * @param code 左子节点0，右子节点1
     * @param sb 用于拼接路径
    */
    public static void getCodes(NodePlus node, String code, StringBuilder sb) {
        StringBuilder stringBuilder2 = new StringBuilder(sb);   //不影响原始数据
        stringBuilder2.append(code);    //往左边拼接0，往右侧拼接1
        if (node != null) {
            //判断当前节点是否是叶子节点
            if (node.getData() == null) {
                //递归处理
                getCodes(node.getLeft(), "0", stringBuilder2);
                getCodes(node.getRight(), "1", stringBuilder2);
            } else {
                huffmanCodes.put(node.getData(), stringBuilder2.toString());
            }
        }
    }

    //统计字符出现的次数
    public static List<NodePlus> getNodes(byte[] bytes) {
        ArrayList<NodePlus> nodes = new ArrayList<NodePlus>();

        //遍历数组，统计每个字符出现的次数
        HashMap<Byte, Integer> countsHashMap = new HashMap<>();
        for (byte b: bytes
             ) {
            Integer count = countsHashMap.get(b);
            if (count == null) {
                countsHashMap.put(b, 1);
            } else {
                countsHashMap.put(b, count + 1);
            }
        }

        //将键值对转换成NodePlus添加到nodes中
        for (Map.Entry<Byte, Integer> entry: countsHashMap.entrySet()
             ) {
            nodes.add(
                    new NodePlus(entry.getKey(), entry.getValue())
            );
        }

        return nodes;
    }

    //创建赫夫曼树
    public static NodePlus huffmanCoding(List<NodePlus> nodes) {
        while (nodes.size() > 1) {
            //1.排序
            Collections.sort(nodes);

            //2.取出最小的两个二叉树，作为左右子树
            NodePlus left = nodes.get(0);
            NodePlus right = nodes.get(1);
            //新的二叉树根节点没有data只有权值
            NodePlus parent = new NodePlus(null, left.getWeight() + right.getWeight());
            parent.setLeft(left);
            parent.setRight(right);

            //
            nodes.remove(left);
            nodes.remove(right);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    public static void preOrder(NodePlus root) {
        if (root != null) {
            root.preOrder();
        } else {
            System.out.println("空树");
        }
    }
}

class NodePlus implements Comparable<NodePlus>{
    private Byte data;  // 存放字符
    private int weight; // 存放权值，字符出现次数
    private NodePlus left;
    private NodePlus right;

    public NodePlus(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    public Byte getData() {
        return data;
    }

    public void setData(Byte data) {
        this.data = data;
    }

    public int getWeight() {
        return weight;
    }

    public void setWeight(int weight) {
        this.weight = weight;
    }

    public NodePlus getLeft() {
        return left;
    }

    public void setLeft(NodePlus left) {
        this.left = left;
    }

    public NodePlus getRight() {
        return right;
    }

    public void setRight(NodePlus right) {
        this.right = right;
    }

    @Override
    public String toString() {
        return "NodePlus{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    @Override
    public int compareTo(NodePlus o) {
        // 从小到大排序
        //当前对象应该排在参数对象之前，则返回一个小于0的整数。
        //两个对象相等（即它们之间没有排序差异），则返回0。
        //当前对象应该排在参数对象之后，则返回一个大于0的整数。
        return this.weight - o.weight;
    }

    public void preOrder() {
//        System.out.print(this.getWeight() + " ");
        System.out.println(this);
        if (this.left != null) {
            this.left.preOrder();
        }
        if (this.right != null) {
            this.right.preOrder();
        }
    }
}